import pandas as pd
import numpy as np
import yfinance as yf
from datetime import datetime, timedelta

class TurtleTradingSystem:
    def __init__(self, capital=100000):
        self.capital = capital
        self.positions = {}
        self.trades = []
        
    def calculate_atr(self, data, period=20):
        """计算真实波动幅度均值(ATR)"""
        high_low = data['High'] - data['Low']
        high_close = np.abs(data['High'] - data['Close'].shift())
        low_close = np.abs(data['Low'] - data['Close'].shift())
        
        true_range = np.maximum(high_low, np.maximum(high_close, low_close))
        atr = true_range.rolling(window=period).mean()
        return atr
    
    def calculate_breakouts(self, data):
        """计算突破点"""
        data['20_day_high'] = data['High'].rolling(window=20).max()
        data['20_day_low'] = data['Low'].rolling(window=20).min()
        data['55_day_high'] = data['High'].rolling(window=55).max()
        data['55_day_low'] = data['Low'].rolling(window=55).min()
        
        data['10_day_exit_high'] = data['High'].rolling(window=10).max()
        data['10_day_exit_low'] = data['Low'].rolling(window=10).min()
        data['20_day_exit_high'] = data['High'].rolling(window=20).max()
        data['20_day_exit_low'] = data['Low'].rolling(window=20).min()
        
        return data
    
    def calculate_position_size(self, price, atr, risk_per_trade=0.01):
        """计算头寸规模"""
        # 每笔交易风险为账户净值的1%
        dollar_risk = self.capital * risk_per_trade
        # 使用2倍ATR作为止损距离
        stop_distance = atr * 2
        # 计算合约/股票数量
        position_size = dollar_risk / stop_distance
        return int(position_size)
    
    def should_enter_long(self, data, i):
        """判断是否应该做多"""
        if i < 55:
            return False
            
        current_close = data['Close'].iloc[i]
        system1_breakout = data['20_day_high'].iloc[i-1]
        system2_breakout = data['55_day_high'].iloc[i-1]
        
        # 系统1或系统2突破
        if current_close > system1_breakout or current_close > system2_breakout:
            return True
        return False
    
    def should_exit_long(self, data, i, system=1):
        """判断是否应该退出多头"""
        if i < 20:
            return False
            
        current_close = data['Close'].iloc[i]
        
        if system == 1:
            exit_level = data['10_day_exit_low'].iloc[i-1]
        else:
            exit_level = data['20_day_exit_low'].iloc[i-1]
            
        return current_close < exit_level
    
    def should_enter_short(self, data, i):
        """判断是否应该做空"""
        if i < 55:
            return False
            
        current_close = data['Close'].iloc[i]
        system1_breakout = data['20_day_low'].iloc[i-1]
        system2_breakout = data['55_day_low'].iloc[i-1]
        
        # 系统1或系统2突破
        if current_close < system1_breakout or current_close < system2_breakout:
            return True
        return False
    
    def should_exit_short(self, data, i, system=1):
        """判断是否应该退出空头"""
        if i < 20:
            return False
            
        current_close = data['Close'].iloc[i]
        
        if system == 1:
            exit_level = data['10_day_exit_high'].iloc[i-1]
        else:
            exit_level = data['20_day_exit_high'].iloc[i-1]
            
        return current_close > exit_level
    
    def backtest(self, symbol, start_date, end_date):
        """回测策略"""
        # 获取数据
        data = yf.download(symbol, start=start_date, end=end_date)
        
        # 计算技术指标
        data['ATR'] = self.calculate_atr(data)
        data = self.calculate_breakouts(data)
        
        # 初始化变量
        cash = self.capital
        position = 0
        entry_price = 0
        system = 1  # 默认使用系统1
        
        for i in range(55, len(data)):
            current_price = data['Close'].iloc[i]
            atr = data['ATR'].iloc[i]
            
            # 检查退出条件
            if position > 0:  # 多头头寸
                if self.should_exit_long(data, i, system):
                    # 平仓
                    cash += position * current_price
                    self.trades.append({
                        'date': data.index[i],
                        'action': 'SELL',
                        'price': current_price,
                        'shares': position,
                        'system': system
                    })
                    position = 0
                    
            elif position < 0:  # 空头头寸
                if self.should_exit_short(data, i, system):
                    # 平仓
                    cash += position * current_price
                    self.trades.append({
                        'date': data.index[i],
                        'action': 'COVER',
                        'price': current_price,
                        'shares': abs(position),
                        'system': system
                    })
                    position = 0
            
            # 检查入场条件
            if position == 0:
                position_size = self.calculate_position_size(current_price, atr)
                
                if self.should_enter_long(data, i):
                    # 开多头
                    if cash >= position_size * current_price:
                        position = position_size
                        cash -= position * current_price
                        entry_price = current_price
                        system = 1 if current_price > data['20_day_high'].iloc[i-1] else 2
                        self.trades.append({
                            'date': data.index[i],
                            'action': 'BUY',
                            'price': current_price,
                            'shares': position,
                            'system': system
                        })
                        
                elif self.should_enter_short(data, i):
                    # 开空头
                    if cash >= position_size * current_price:
                        position = -position_size
                        cash -= position * current_price  # 做空是借券卖出，获得现金
                        entry_price = current_price
                        system = 1 if current_price < data['20_day_low'].iloc[i-1] else 2
                        self.trades.append({
                            'date': data.index[i],
                            'action': 'SHORT',
                            'price': current_price,
                            'shares': abs(position),
                            'system': system
                        })
            
            # 更新账户净值
            self.capital = cash + position * current_price
        
        return data, self.trades

# 使用示例
if __name__ == "__main__":
    # 初始化海龟交易系统
    turtle = TurtleTradingSystem(capital=100000)
    
    # 运行回测
    symbol = "AAPL"
    start_date = "2020-01-01"
    end_date = "2024-01-01"
    
    data, trades = turtle.backtest(symbol, start_date, end_date)
    
    # 打印交易记录
    print(f"初始资金: $100,000")
    print(f"最终资金: ${turtle.capital:,.2f}")
    print(f"总收益率: {(turtle.capital - 100000) / 100000 * 100:.2f}%")
    print(f"交易次数: {len(trades)}")
    
    # 显示前10笔交易
    print("\n前10笔交易:")
    for trade in trades[:10]:
        print(f"{trade['date'].strftime('%Y-%m-%d')} {trade['action']} {trade['shares']}股 @ ${trade['price']:.2f}")